import matplotlib.pyplot as plt  # 加载matplotlib用于数据的可视化
from sklearn.decomposition import PCA  # 加载PCA算法包

class PCAAnalysis:
    def __init__(self, n_components=2):
        '''
        n_components: 降维后主成分数目
        '''
        self.n_components = n_components
        self.pca = PCA(n_components=n_components)
        self.reduced_x = None
        self.y = None
        self.x = None

    def fit_transform(self, x, y=None):
        '''
        x: 原始数据，shape=(n_samples, n_features)
        y: 标签，可选，shape=(n_samples,)
        '''
        self.x = x
        self.y = y
        self.reduced_x = self.pca.fit_transform(x) #形状是(n_samples, n_components)
        return self.reduced_x

    def plot(self, labels=None, colors=None, markers=None):
        '''
        labels: 标签数组，shape=(n_samples,)
        colors: 颜色列表，如['r','b','g']
        markers: 点型列表，如['x','D','.']
        '''
        if self.reduced_x is None:
            raise ValueError('请先调用fit_transform方法')
        if labels is None:
            plt.scatter(self.reduced_x[:,0], self.reduced_x[:,1], c='gray', marker='o')
        else:
            unique_labels = sorted(set(labels))
            if colors is None:
                colors = ['r','b','g','c','m','y','k']
            if markers is None:
                markers = ['o','x','D','.','^','s','*']
            for i, lab in enumerate(unique_labels):
                idx = [j for j in range(len(labels)) if labels[j]==lab]
                plt.scatter(self.reduced_x[idx,0], self.reduced_x[idx,1], c=colors[i%len(colors)], marker=markers[i%len(markers)], label=str(lab))
            plt.legend()
        plt.xlabel('PC1')
        plt.ylabel('PC2')
        plt.title('PCA降维结果')
        plt.show()

# 示例：使用鸢尾花数据集
if __name__ == "__main__":
    from sklearn.datasets import load_iris
    data = load_iris()
    x = data.data
    y = data.target
    #x是原始数据，y是分类的标签“花1，花2”
    pca_model = PCAAnalysis(n_components=2)
    pca_model.fit_transform(x, y)
    pca_model.plot(labels=y)